from math import sqrt
from numpy import dot,zeros_like,array
from v3dsolver.neighbour_finder import Bin, NeighbourFinder
from v3dsolver.fluid_elements.fluid_elems import Saffman
from v3dsolver.vector import Vec
import pdb

class MergingManager:
    """
    This class is responsible for merging vortices which come very close to 
    each other into a single vortex having their combined vorticity.
    """
    def __init__(self,elems,nu,dt,R=sqrt(8),rmerge_cellsize=0.4,
                   cutoff_fraction=0.001):
        """
        Initialization of merging manager.
        
        elems: elements stored in current processor
        nu: kinematic viscosity of the fluid
        dt: timestep of the simulation
        R: radius for binning (sqrt(8) for 3D flow; sqrt(12) for 2D flow)
        rmerge_cellsize: radius within which particles will be merged as a 
                       fraction of cellsize (default: 0.4 (from Arnav's thesis))
        cutoff_fraction: fraction of avg vorticity below which a particle will
                         be discarded
        """
        self.elems = elems
        self.nu = nu
        self.dt = dt
        self.R = R
        self.cellsize = R*sqrt(self.nu*self.dt)
        self.r_merge = rmerge_cellsize*self.cellsize # r_merge: merging radius
        self.cutoff_fraction = cutoff_fraction
    
    def merge(self):
        bin_obj = Bin(self.elems,self.cellsize)
        nbrfinder = NeighbourFinder(bin_obj,self.nu,self.dt)
        merged = zeros_like(self.elems) # 0 if elem is not merged, 1 o/w
        new_elems = []
        for elemindex,elem in enumerate(self.elems):
            nbrs = nbrfinder.get_neighbours_serial(elem)
            for nbr in nbrs: # nbr is of type [nbr_hash,nbr_pos]
                nbrindex = int(nbr[0].split()[-1])
                nbrelem = self.elems[nbrindex]
                r = (elem.pos - nbrelem.pos)
                r = r.length()
                if (r<self.r_merge and merged[elemindex]==0 
                    and nbrindex!=elemindex):
                    merged[elemindex] = 1
                    merged[nbrindex] = 1
                    elemvor = elem.vor.length()
                    nbrvor = nbrelem.vor.length()
                    if (elemvor+nbrvor)>1e-5:
                        new_pos = ((elem.pos*elemvor + nbrelem.pos*nbrvor)/
                                   (elemvor + nbrvor))
                        new_vor = elem.vor + nbrelem.vor
                        new_elems += [Saffman(pos=new_pos,vel=array([0.,0.,0.]),
                                      vor=new_vor,rad=self.elems[0].rad)]
        
        # update self.elems according to new merged particles
        for i,elem in enumerate(self.elems[:]):
            if merged[i]==1: # if merged somewhere
                self.elems.remove(elem)
        self.elems += new_elems
    
    def eat(self):
        """
        Eats up elements having vorticity less than the cutoff fraction of the
        average vorticity.
        """
        avg_vor = Vec()
        for elem in self.elems:
            avg_vor += elem.vor
        avg_vor = avg_vor/(float(len(self.elems)))
        avg_vor_mag = avg_vor.length() # magnitude of avg_vor
        
        cutoff_vor = self.cutoff_fraction*avg_vor_mag
        for elem in self.elems:
            vor = elem.vor.length()
            if vor<=cutoff_vor:
                self.elems.remove(elem)
